package org.dromara.common.websocket.holder;

import lombok.AccessLevel;
import lombok.NoArgsConstructor;
import org.springframework.web.socket.WebSocketSession;

import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CopyOnWriteArrayList;

/**
 * WebSocketSession 用于保存当前所有在线的会话信息
 *
 * @author zendwang
 */
@NoArgsConstructor(access = AccessLevel.PRIVATE)
public class WebSocketSessionHolder {

    private static final Map<String, List<WebSocketSession>> USER_SESSION_MAP = new ConcurrentHashMap<>();

    public static void addSession(String sessionKey, WebSocketSession session) {
        USER_SESSION_MAP.computeIfAbsent(sessionKey, k -> new CopyOnWriteArrayList<>())
                .add(session);
    }

    public static void removeSession(String sessionKey, String sessionId) {
        USER_SESSION_MAP.computeIfPresent(sessionKey, (key, sessions) -> {
            // 使用并发安全的移除操作
            boolean removed = sessions.removeIf(session -> sessionId.equals(session.getId()));

            // 自动清理空列表，防止内存泄漏
            return removed && sessions.isEmpty() ? null : sessions;
        });
    }

    public static List<WebSocketSession> getSessions(String sessionKey) {
        return USER_SESSION_MAP.get(sessionKey);
    }

    public static Set<String> getSessionsAll() {
        return USER_SESSION_MAP.keySet();
    }

    public static Boolean existSession(String sessionKey) {
        return USER_SESSION_MAP.containsKey(sessionKey);
    }
}
